import os

import numpy as np
import xarray as xr
# import numpy as np

"""
2023era5 = r'G:\GEP_data\era5\2023era5data'
2024era5 = r'G:\GEP_data\era5\2024era5data'
"""

# 读取蒸散发和气温数据
# input_dir = r'G:\GEP_data\era5\2023era5data'
# year = 2023
# output_dir = os.path.join(input_dir, 'output')
# os.makedirs(output_dir, exist_ok=True)
# temperature_file = os.path.join(input_dir, '2023_temp_hour.nc')
# evaporation_file = os.path.join(input_dir, 'Evaporation_hour_2023.nc')


input_dir = r'G:\GEP_data\era5\2023era5data'
year = 2023
output_dir = os.path.join(input_dir, 'output')
os.makedirs(output_dir, exist_ok=True)
temperature_file = os.path.join(input_dir, '2023_temp_hour.nc')
evaporation_file = os.path.join(input_dir, 'Evaporation_hour_2023.nc')


nodata = -9999
evaporation_dataset = xr.open_dataset(evaporation_file)
temperature_dataset = xr.open_dataset(temperature_file)

# 确保时间维度对齐
evaporation_dataset = evaporation_dataset.sel(valid_time=temperature_dataset['valid_time'])

# 筛选气温大于26°C的时间段
mask = temperature_dataset[list(temperature_dataset.data_vars.keys())[0]] > 26.0 + 273.15

# 计算开放空调降温期间的蒸散发量
E_wt = evaporation_dataset[list(evaporation_dataset.data_vars.keys())[0]].where(mask, drop=True)

annual_e_wt = E_wt.resample(valid_time='1YE').sum(dim='valid_time')

# 统计气温大于26°C的时间段数量
num_hours_gt_26 = mask.sum(dim='valid_time').values

# 将小时数转换为天数，并向上取整
num_days_gt_26 = np.ceil(num_hours_gt_26 / 24)
num_days_gt_26[num_days_gt_26 == 0] = 1

# 保存结果为 NetCDF 文件
output_nc_file = os.path.join(output_dir, f'Ewt_{year}.nc')
annual_e_wt.to_netcdf(output_nc_file)

# 保存结果为 TIFF 文件
output_tif_file = os.path.join(output_dir, f'Ewt_{year}.tif')

# 确保NetCDF文件中的地理坐标系统和投影方式被正确设置
# 假设NetCDF文件中包含地理坐标信息
annual_e_wt_riox = annual_e_wt.rio.write_crs("EPSG:4326")  # 设置为WGS84坐标系统
annual_e_wt_riox = annual_e_wt_riox.rio.write_nodata(nodata)  # 设置无效值为NaN
annual_e_wt_riox = annual_e_wt_riox.astype(np.float32)  # 转换为32位浮点数
annual_e_wt_riox.rio.to_raster(output_tif_file)

# 将num_days_gt_26保存为TIFF文件
output_days_tif_file = os.path.join(output_dir, f'ETT-D_{year}.tif')

# 创建一个新的xarray DataArray来保存num_days_gt_26
num_days_gt_26_da = xr.DataArray(num_days_gt_26, dims=['latitude', 'longitude'],
                                 coords={'latitude': annual_e_wt.latitude,
                                         'longitude': annual_e_wt.longitude})

# 将DataArray转换为GeoTIFF
num_days_gt_26_da_riox = num_days_gt_26_da.rio.write_crs("EPSG:4326")  # 设置为WGS84坐标系统
num_days_gt_26_da_riox = num_days_gt_26_da_riox.rio.write_nodata(nodata)  # 设置无效值为NaN
num_days_gt_26_da_riox.rio.to_raster(output_days_tif_file)

# 输出结果
print(f"数据已保存到 {output_nc_file}")
print(f"数据已保存到 {output_tif_file}")
print(f"数据已保存到 {output_days_tif_file}")

# 关闭数据文件
evaporation_dataset.close()
temperature_dataset.close()